package cn.getech.data.development.sink.kudu

import java.io.{ByteArrayOutputStream, ObjectOutputStream}
import java.lang

import cn.getech.data.development.sink.jdbc.config.CustomTableConfig
import cn.getech.data.development.utils.{ConfigurationManager, JDBCSQLTypeUtils}
import org.apache.flink.api.java.tuple
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction
import org.apache.flink.types.Row
import org.apache.kudu.client._
import org.apache.kudu.{ColumnSchema, Schema, Type}
import org.slf4j.{Logger, LoggerFactory}

// todo test
class FlinkStreamKuduUpsertSink(conf: CustomTableConfig) extends RichSinkFunction[tuple.Tuple2[lang.Boolean, Row]] {
  private val logger: Logger = LoggerFactory.getLogger(this.getClass)
  private var client: KuduClient = null
  private var table: KuduTable = null
  private var kuduSession: KuduSession = null
  private var out: ByteArrayOutputStream = null
  private var os: ObjectOutputStream = null

  override def open(parameters: Configuration): Unit = {
    val config = conf.getConfig()
    out = new ByteArrayOutputStream
    os = new ObjectOutputStream(out)
    client = new KuduClient.KuduClientBuilder(s"${config.url._2}").build
    if (client.tableExists(config.tablename._2)) {
      table = client.openTable(config.tablename._2)
    }
    kuduSession = client.newSession
    kuduSession.setFlushMode(SessionConfiguration.FlushMode.MANUAL_FLUSH)
    kuduSession.setMutationBufferSpace(conf.getBatchInterval)
  }

  override def invoke(value: tuple.Tuple2[lang.Boolean, Row]): Unit = {
    var row: PartialRow = null
    var operator: Operation = null
    if (value.f0) operator = table.newInsert
    else operator = table.newDelete()
    row = operator.getRow
    val utils = new JDBCSQLTypeUtils
    for (i <- 0 until conf.getParamSize) {
      utils.kudu_type2DataType(row, conf.getParams(i), value.f1.getField(i).toString)
    }
    kuduSession.flush
    val response = kuduSession.apply(operator)
    if (response != null) logger.error(response.getRowError.toString)
  }

  override def close(): Unit = {
    try {
      kuduSession.close
      client.close()
      os.close()
      out.close
    } catch {
      case e: Exception =>
        logger.error(e.getMessage)
    }
  }

  private def kuduCreate(): Unit = {
    // 创建kudu连接
    val kuduClient = new KuduClient.KuduClientBuilder(ConfigurationManager.getProperty("kudu.master")).build()
    // 设置表名
    val tableName = "TEST_PERFORMANCE_KUDU"
    // 创建列
    val colums = List[ColumnSchema]((new ColumnSchema.ColumnSchemaBuilder("userid", Type.STRING).key(true).nullable(false).build()),
      (new ColumnSchema.ColumnSchemaBuilder("action", Type.STRING).nullable(true).build()),
      (new ColumnSchema.ColumnSchemaBuilder("duration", Type.INT64).nullable(true).build()),
      (new ColumnSchema.ColumnSchemaBuilder("dt", Type.INT64).nullable(true).build()))

    import scala.collection.JavaConverters._
    val schema: Schema = new Schema(colums.asJava)

    val cto: CreateTableOptions = new CreateTableOptions()
    cto.setRangePartitionColumns(List("userid").asJava).setNumReplicas(1)
    kuduClient.createTable(tableName, schema, cto)
    kuduClient.close()
  }
}
